# loss_manager.py
from typing import List, Any, Dict

import torch
from torch import nn

from config_manager import LossConfig


class LossManager:
    """
    根据配置构建多个损失函数，并自动计算加权总损失。
    支持：
      - cross_entropy
      - mse
      - l1
      - bce
      - huber
      - 自定义 Callable
    """
    def __init__(self, loss_config: LossConfig):
        self.loss_config = loss_config
        self.loss_fns = self._build_losses()

    # 构建损失函数
    def _build_losses(self) -> List[dict]:
        """
        构建损失函数方法
        """
        loss_fns = []
        for cfg in self.loss_config.losses:
            if cfg.custom_loss_fn is not None:
                fn = cfg.custom_loss_fn
            elif cfg.loss_type == "cross_entropy":
                # 统计正确的类别概率占比
                fn = nn.CrossEntropyLoss(
                    label_smoothing=cfg.label_smoothing,
                    reduction=cfg.reduction
                )
            elif cfg.loss_type == "mse":
                # 均方误差损失 回归任务，衡量模型预测值与真实值之间的平方差 比较对应值的差距
                fn = nn.MSELoss(reduction=cfg.reduction)
            elif cfg.loss_type == "l1":
                # 平均绝对误差损失 计算预测值与真实值之间的 绝对差的平均值
                fn = nn.L1Loss(reduction=cfg.reduction)
            elif cfg.loss_type == "bce":
                # 二分类交叉熵损失
                fn = nn.BCELoss(reduction=cfg.reduction)
            elif cfg.loss_type == "bce_with_logits":
                # 二分类交叉熵损失（BCELoss） 的升级版本，它结合了：
                # Sigmoid 激活：将 logits 转为概率
                # Binary Cross Entropy：计算概率与真实标签之间的差异
                fn = nn.BCEWithLogitsLoss(reduction=cfg.reduction)
            elif cfg.loss_type == "huber":
                # 平滑 L1 损失
                fn = nn.HuberLoss(reduction=cfg.reduction)
            else:
                raise ValueError(f"❌ 未知的损失类型: {cfg.loss_type}")

            loss_fns.append({
                "name": cfg.loss_type or "custom",
                "fn": fn,
                "weight": cfg.weight
            })
        return loss_fns

    # 计算总损失
    def compute_loss(self, predictions: Any, targets: Any) -> Dict[str, torch.Tensor]:
        """
        根据配置计算加权总损失。
        支持：
            - 单任务: predictions, targets 都是 Tensor
            - 多任务: predictions, targets 为字典，key 对应 loss_config 顺序
        返回:
            {
                "total_loss": 总损失,
                "details": {loss_name: 单项损失值}
            }
        """
        total_loss = 0.0
        detail_losses = {}

        # 多任务情况下传入字典
        is_multi_task = isinstance(predictions, dict)

        for i, loss_item in enumerate(self.loss_fns):
            fn = loss_item["fn"]
            weight = loss_item["weight"]
            name = loss_item["name"]

            # 获取当前任务的预测与标签
            if is_multi_task:
                pred = predictions.get(name, None)
                target = targets.get(name, None)
                if pred is None or target is None:
                    raise KeyError(f"❌ Multi-task 模式下缺少键: {name}")
            else:
                pred, target = predictions, targets

            loss_val = fn(pred, target) * weight
            total_loss += loss_val
            detail_losses[name] = loss_val.detach()

        return {"total_loss": total_loss, "details": detail_losses}

